# https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from PyCmpltrtok.common import sep

# On Windows platform, the torch.distributed package only
# supports Gloo backend, FileStore and TcpStore.
# For FileStore, set init_method parameter in init_process_group
# to a local file. Example as follow:
# init_method="file:///f:/libtmp/some_file"
# dist.init_process_group(
#    "gloo",
#    rank=rank,
#    init_method=init_method,
#    world_size=world_size)
# For TcpStore, same way as on Linux.


def setup(rank, world_size):
    # dist.init_process_group("nccl", init_method='env://')
    dist.init_process_group("gloo", init_method='env://')


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic():
    local_rank, rank, world_size = [int(x) for x in (os.environ['LOCAL_RANK'], os.environ['RANK'], os.environ['WORLD_SIZE'], )]
    flag = f'{local_rank}/{rank}/{world_size}'
    sep(f'{flag} start')
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)
    print(flag, 'after setup')

    # create model and move it to GPU with id rank
    print(flag, 'before build model')
    model = ToyModel().to(local_rank)
    ddp_model = DDP(model, device_ids=[local_rank])
    print(flag, 'after build model')

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    print(flag, '20x10 ->', outputs.shape)
    labels = torch.randn(20, 5).to(local_rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    sep(f'{flag} start cleanup')
    cleanup()
    sep(f'{flag} cleanup end')


if '__main__' == __name__:

    demo_basic()
